#include <butil/string_printf.h>

#include "index/primary_key_index/primary_key_disk_indexer.h"

namespace hawking {
namespace indexlib {

DECLARE_int32(primary_key_default_pool_bytes);

Status PrimaryKeyDiskIndexer::Load(
    const std::string& pk_dir,
    std::shared_ptr<DeletionMapIndexer> deletion_indexer) {
    std::string data_file = butil::string_printf("%s/data_file", pk_dir.data());
    pk_disk_reader_ = std::make_unique<MmapFileReader>(data_file);
    Status ret = pk_disk_reader_->Init();
    if (ret != Status::OK) {
        LOG(ERROR) << "primary_key_reader Init " << data_file
            << " Fail with " << static_cast<int32_t>(ret);
        return ret;
    }

    pk_disk_reader_->Read((uint8_t*)&min_pk_, sizeof(min_pk_), 0).Value();
    pk_disk_reader_->Read(
        (uint8_t*)&max_pk_,
        sizeof(max_pk_),
        pk_disk_reader_->Length() - pk_disk_item_length_).Value();

    return Status::OK;
}

void PrimaryKeyDiskIndexer::ExcludeTombPks(
    const HashSet<uint64_t>& tomb_stone_pks,
    std::shared_ptr<DeletionMapIndexer> deletion_indexer) {
    HashSet::Iterator iter = tomb_stone_pks.CreateIterator();
    while (iter.HasNext()) {
        uint64_t deleted_pk = iter.Next();
        ReturnValue<DocId> del_doc = Lookup(deleted_pk);
        if (del_doc.Code() == Status::OK) {
            deletion_indexer->Delete(del_doc.Value());
        }
    }
}

bool PrimaryKeyDiskIndexer::IsExist(uint64_t pk) const {
    if (pk < min_pk_ || pk > max_pk_) {
        return false;
    }

    size_t low = 0;
    size_t high = pk_disk_reader_->Length() / pk_disk_item_length_ - 1;

    while (low <= high) {
        size_t mid = low + ((high - low) >> 1);
        uint64_t mid_pk;
        pk_disk_reader_->Read(
            (uint8_t*)&mid_pk, sizeof(mid_pk), mid * pk_disk_item_length_);

        if (mid_pk == pk) {
            return true;
        } else if (mid_pk < pk) {
            low = mid + 1;
        } else {
            high = mid - 1;
        }
    }

    return false;
}

void PrimaryKeyDiskIndexer::LookupOne_(uint64_t pk, SegmentContext* segment_ctx) {
    if (pk < min_pk_ || pk > max_pk_) {
        return;
    }

    size_t low = 0;
    size_t high = pk_disk_reader_->Length() / pk_disk_item_length_ - 1;

    while (low <= high) {
        size_t mid = low + ((high - low) >> 1);
        uint64_t mid_pk;
        pk_disk_reader_->Read(
            (uint8_t*)&mid_pk, sizeof(mid_pk), mid * pk_disk_item_length_);

        if (mid_pk == pk) {
            DocId doc_id;
            pk_disk_reader_->Read(
                (uint8_t*)&doc_id,
                sizeof(doc_id),
                mid * pk_disk_item_length_ + pk_length_);
            segment_ctx->PushDocId(doc_id);
            segment_ctx->MutableDocContext()->back().pk = pk;
            return;
        } else if (mid_pk < pk) {
            low = mid + 1;
        } else {
            high = mid - 1;
        }
    }
}

Status PrimaryKeyDiskIndexer::Lookup(
    const google::protobuf::RepeatedField<int64_t>& pks,
    SegmentContext* segment_ctx) {
    segment_ctx->Reserve(pks.size());
    for (const auto pk : pks) {
        LookupOne_(static_cast<uint64_t>(pk), segment_ctx);
    }

    return Status::OK;
}

ReturnValue<DocId> PrimaryKeyDiskIndexer::Lookup(uint64_t pk) {
    if (pk < min_pk_ || pk > max_pk_) {
        return ReturnValue<DocId>{Status::NotFound, 0};
    }

    size_t low = 0;
    size_t high = pk_disk_reader_->Length() / pk_disk_item_length_ - 1;

    while (low <= high) {
        size_t mid = low + ((high - low) >> 1);
        uint64_t mid_pk;
        pk_disk_reader_->Read(
            (uint8_t*)&mid_pk, sizeof(mid_pk), mid * pk_disk_item_length_);

        if (mid_pk == pk) {
            DocId doc_id;
            pk_disk_reader_->Read(
                (uint8_t*)&doc_id,
                sizeof(doc_id),
                mid * pk_disk_item_length_ + pk_length_);
            return ReturnValue<DocId>{Status::OK, doc_id};
        } else if (mid_pk < pk) {
            low = mid + 1;
        } else {
            high = mid - 1;
        }
    }

    return ReturnValue<DocId>{Status::NotFound, 0};
}

}
}